import cv2
import numpy as np


# 遮盖下半部分区域
def lower_mask(image):
    height = image.shape[0]
    width = image.shape[1]

    # 定义一个四边形区域
    mask = np.zeros_like(image)
    polygon = np.array([[
        (0, 0),
        (0, height/2),
        (width, height/2),
        (width, 0)
    ]], np.int32)

    cv2.fillPoly(mask, polygon, (255,255,255))
    masked_image = cv2.bitwise_and(image, mask)
    return masked_image

# 提取图像蓝色部分
def extract_blue_region(image):
    # 将 BGR 图像转换为 HSV 图像
    hsv_img = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    # 设置蓝色的 HSV 范围
    lower_blue = np.array([100, 150, 100])   # 蓝色的低范围
    upper_blue = np.array([130, 255, 200])  # 蓝色的高范围

    # 使用 inRange() 函数提取蓝色区域
    blue_mask = cv2.inRange(hsv_img, lower_blue, upper_blue)

    # 根据掩膜提取蓝色区域
    blue_region = cv2.bitwise_and(image, image, mask=blue_mask)

    return blue_region


if __name__ == "__main__":
    image_path = "dataset/image5.jpg"
    image = cv2.imread(image_path)

    masked_img = lower_mask(image)
    blue_region = extract_blue_region(masked_img)

    cv2.imshow('masked_region', masked_img)
    cv2.imshow('blue_region', blue_region)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    # 高斯模糊
    blurred = cv2.blur(blue_region, (9, 9))
    cv2.imshow('blurred', blurred)

    # 二值化
    ret, binary = cv2.threshold(blurred, 50, 255, cv2.THRESH_BINARY)
    cv2.imshow('blurred binary', binary)

    # 做闭运算消除内部空隙
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (21, 7))
    closed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
    cv2.imshow('closed', closed)

    cv2.waitKey(0)
    cv2.destroyAllWindows()


    # 绘制外接矩形
    gray_closed = cv2.cvtColor(closed, cv2.COLOR_BGR2GRAY)
    contours, _ = cv2.findContours(gray_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 遍历所有轮廓
    for contour in contours:
        # 计算最小外接矩形
        rect = cv2.minAreaRect(contour)
        # 获取矩形的角点
        box = cv2.boxPoints(rect)
        box = np.array(box, dtype=np.int32)  # 转换为整数

        # 在原图像上绘制最小外接矩形
        cv2.drawContours(image, [box], 0, (0, 255, 0), 3)

    cv2.imshow('signal_detect', image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
